[WIP] [feat] delta embedding dump upload to PAI Feature Store#603
[WIP] [feat] delta embedding dump upload to PAI Feature Store#603eric-gecheng wants to merge 14 commits into
Conversation
| # Step zero is excluded from the delta publication contract. | ||
| logger.info("Skipping delta embedding dump at step %s.", global_step) | ||
| return None | ||
| if self._interval_steps is not None and global_step % self._interval_steps == 0: |
There was a problem hiding this comment.
After _sync_final_step() replaces each local step with the global MAX, this boundary may only have been dumped by the longest-running rank. For example, if ranks finish at steps 49 and 50, the shorter rank adopts 50 and skips here, so its trailing tracked rows are never written. Please reject uneven exhaustion for all multi-rank dump cadences, or base this skip on each rank’s actual last dump while preserving collective participation; add a 49/50 regression test.
| should_dump = self._rank == 0 and time.monotonic() >= self._next_dump_time | ||
|
|
||
| any_failed = local_error is not None | ||
| if getattr(self, "_world_size", 1) > 1 and self._interval_secs is not None: |
There was a problem hiding this comment.
This failure rendezvous is limited to minute cadence. In default step mode, rank zero can observe an async uploader failure immediately while peers skip the throttled marker check, raise alone, and leave peers hanging in the next training collective. Please synchronize the failure bit for every distributed FeatureStore cadence before raising; boundary dump failures need the same treatment.
| dtype=torch.int32, | ||
| device=self._device, | ||
| ) | ||
| torch.distributed.all_reduce(state, op=torch.distributed.ReduceOp.MAX) |
There was a problem hiding this comment.
Minute cadence executes this CUDA all_reduce plus host-synchronizing .item() calls every training step, while the enabled pipeline also all-reduces batch availability. That adds at least two blocking collectives per batch for dumps minutes apart. Please coalesce timer/error state with the existing availability rendezvous or otherwise remove redundant hot-path synchronization.
| table = pq.read_table(source) | ||
| source.seek(0) | ||
| self._validate_parquet_schema(table.schema, path) | ||
| for row in table.to_pylist(): |
There was a problem hiding this comment.
read_table() + to_pylist() materializes every embedding element as Python objects, then copies rows into NumPy arrays and another sorted list before uploading. One large delta can therefore exhaust rank-zero training memory despite max_pending_steps. Please stream Parquet batches and use bounded/spilled deduplication so memory is capped independently of one dump’s size.
| with self._condition: | ||
| if self._aborting: | ||
| return | ||
| self._add_discovered_steps_locked() |
There was a problem hiding this comment.
The worker calls _add_discovered_steps_locked() on every poll. Because committed dumps/manifests are retained, each poll rescans all historical steps, rereads success/manifest JSON and Parquet schemas, and does so while holding _condition, so submit()/close() latency grows with job history. Please reconcile once at startup and incrementally track later submissions/new entries.
| single_rank_path = os.path.join( | ||
| output_dir, f"{file_prefix}_step_{global_step}.parquet" | ||
| ) | ||
| if os.path.isfile(single_rank_path): |
There was a problem hiding this comment.
These canonical paths are sampled without validating their hashes/generation against the selected success marker. A restored run can replace the same step’s canonical files before the new generation commits, making this tool compare uncommitted local rows with the previous remote commit. Please verify the selected files against success["shards"], or fail explicitly when the committed bytes are no longer available.
| @@ -0,0 +1 @@ | |||
| feature_store_py @ https://gecheng-rec.oss-accelerate-overseas.aliyuncs.com/software/feature_store_py-2.2.7-py3-none-any.whl | |||
There was a problem hiding this comment.
Please source this credential-bearing SDK from an organization-controlled registry/release and pin its SHA-256 hash. A mutable wheel URL in a personal bucket makes dependency compromise equivalent to arbitrary code execution with the training job’s FeatureStore credentials.
| "feature_store_config.region must not be empty " | ||
| "(it may come from ALIBABA_CLOUD_REGION)" | ||
| ) | ||
| endpoint_url = endpoint if "://" in endpoint else f"//{endpoint}" |
There was a problem hiding this comment.
This only rejects URI userinfo; an arbitrary plaintext or attacker-controlled endpoint is still passed to the SDK together with cloud and FeatureDB credentials. Please require HTTPS and trusted FeatureStore endpoints (or an explicit trusted custom-endpoint opt-in), and reject unexpected path/query/fragment components.
| // Existing FeatureStore entity associated with an automatically created | ||
| // DynamicEmbedding FeatureView. | ||
| required string feature_entity_name = 8; | ||
| // Explicit FeatureDB version for this incremental training run. It must be |
There was a problem hiding this comment.
This says the version must be provisioned before upload, but the uploader contract is AUTO_CREATE_ON_FIRST_DELTA_MERGE and the implementation deliberately performs no version precheck. Please document that the explicit non-default version is created by the first MERGE.
|
Static review summary (tests/builds intentionally not run):
Detailed actionable findings are inline. |
No description provided.